import os

os.environ['CUDA_VISIBLE_DEVICES']='0'
import PIL.Image
from PIL import Image,ImageOps
import torch,pdb,PIL,shutil
from diffusers import FluxFillPipeline,FluxPriorReduxPipeline
from diffusers.utils import load_image
from itertools import product

from util_flux import pad_image
from util_flux import horizontal_concat_images

from util_sam import get_mask_by_sam
from util_mask import get_erosed_mask_by_radtio,add_random_holes


FLUX_FILL='/home/shengjie/ckp/FLUX.1-Fill-dev'
FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX_DEPTH='/home/shengjie/ckp/FLUX.1-Depth-dev'
FLUX_DEPTH_LORA='/home/shengjie/ckp/FLUX.1-Depth-dev-lora'
FLUX='/data/models/FLUX___1-dev'

DEPTH_PREDCITION='/home/shengjie/ckp/depth-anything-large-hf'

target_shape = (1024,1024)

# mask,mask_pil = get_erosed_mask_by_radtio(mask,0.5)

## emd of img1_ori
## img of img1_edit
## mask of img1_mask

# emd_of_img_ori_path = '/data/shengjie/style_zhenzhi/img_2.jpg'
# img_of_img_edit_path = '/data/shengjie/style_zhenzhi_ori_edited_mask/img_2.jpg'
# mask_of_img_mask_path = '/data/shengjie/style_zhenzhi_ori_edited_mask/img_2_mask.png'

# save_path = '/data/shengjie/style_zhenzhi/img_2-1.jpg'

# niukou
# niukou-pockets
# yinhua
types = ['niukou','niukou-pockets','yinhua']
choose_type = types[2]

ori_dir = f'/mnt/nas/shengjie/datasets_zhenzhi/zhenzhi_extend'
edit_dir = f'/mnt/nas/shengjie/datasets_zhenzhi/zhenzhi-{choose_type}-edited'
save_dir = f'/mnt/nas/shengjie/datasets_zhenzhi/zhenzhi-{choose_type}-edited-restore'
# if os.path.exists(save_dir):
#     shutil.rmtree(save_dir)
os.makedirs(save_dir,exist_ok=True)

check_name = lambda n :  '-3' == os.path.splitext(n)[0][-2:]


# 破碎的img
# xxxx-1/2/3.jpg
img_of_img_edit_paths = sorted([os.path.join(edit_dir,n) 
                                for n in os.listdir(edit_dir) 
                                if os.path.splitext(n)[1]=='.jpg' and \
                                    check_name(n)])

# 原图作为emb 用于修复图片
# xxxx.jpg
emd_of_img_ori_paths = sorted([ os.path.join(ori_dir,
                                    os.path.basename(n)[:-6]+os.path.splitext(n)[1]) 
                                for n in img_of_img_edit_paths
                                ])

# img mask
# xxxx-1/2/3.png
mask_of_img_mask_paths = sorted([ os.path.splitext(n)[0]+'_mask.png'
                                for n in img_of_img_edit_paths
                                ])
# 用 原图emb 修复 破碎的img
save_paths = [ os.path.join(save_dir , os.path.basename(n) )  
              for n in img_of_img_edit_paths]

# pdb.set_trace()
assert len(emd_of_img_ori_paths) == len(img_of_img_edit_paths),\
    f'{len(emd_of_img_ori_paths)} != {len(img_of_img_edit_paths)}'
# pdb.set_trace()

def process_img_1024(imgpath,pad_color=(255,255,255)) -> PIL.Image:
    image = load_image(imgpath) # PIL
    image,_,_,_,_ = pad_image(image,fill=pad_color)
    image = image.resize(target_shape)
    return image

pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                                FLUX_REDUX, 
                                                torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxFillPipeline.from_pretrained(FLUX_FILL, torch_dtype=torch.bfloat16).to("cuda")
    
extra_re_gennerate = []


def generate(emd_of_img_ori_path,img_of_img_edit_path,mask_of_img_mask_path):
    emd_of_img_ori = process_img_1024(emd_of_img_ori_path)
    img_of_img_edit = process_img_1024(img_of_img_edit_path)
    mask_of_img_mask = process_img_1024(mask_of_img_mask_path,pad_color=(0,0,0))

    assert isinstance(mask_of_img_mask,PIL.Image.Image)
    mask = mask_of_img_mask.convert('L')
    mask_invert = ImageOps.invert(mask)
    # mask_invert.save('tmp.jpg')
    background = PIL.Image.new( 'RGB' , img_of_img_edit.size, (255,255,255) )
    img_of_img_edit = PIL.Image.composite( img_of_img_edit , background , mask_invert )
    # img_of_img_edit.save('tmp.jpg')

    main_condition_prompt = pipe_prior_redux(emd_of_img_ori) # attr 'prompt_embeds' torch.Size([1, 1241, 4096]) 

    res_image = pipe(
                # prompt="green collarless",
                image=img_of_img_edit,
                mask_image=mask_of_img_mask,
                height=target_shape[1],
                width=target_shape[0],
                guidance_scale=30,
                num_inference_steps=8,
                max_sequence_length=512,
                generator=torch.Generator("cpu").manual_seed(0),
                **main_condition_prompt,
            ).images[0]


    concat_res = horizontal_concat_images([img_of_img_edit,emd_of_img_ori,
                                        mask_of_img_mask,
                                        res_image])
    return concat_res,res_image

emd_of_img_ori_path,img_of_img_edit_path,mask_of_img_mask_path = \
'cloth_collar.jpg','cloth.jpg','cloth_mask.png'

concat_res,res_image = generate( emd_of_img_ori_path,img_of_img_edit_path,mask_of_img_mask_path )

concat_res.save('tmp.jpg')

def main():
    for emd_of_img_ori_path,img_of_img_edit_path,mask_of_img_mask_path,save_path in \
    zip(emd_of_img_ori_paths,img_of_img_edit_paths,mask_of_img_mask_paths,save_paths):
        
        
        # if len(extra_re_gennerate)==0:pass
        # else:
        #     is_continue = True
        #     for e in extra_re_gennerate:
        #         if e in emd_of_img_ori_path:
        #             is_continue = False
        #             break
        #     if is_continue:continue
            # import re
            # pattern = re.compile(r'img_\d+')
            # # re.compile(r'img_(\d+(?:-\d+)?)')
            # pdb.set_trace()
            # if pattern.search(emd_of_img_ori_path).group() \
            #     not in extra_re_gennerate:
            #     print('不在extra中，无需处理 continue')
            #     continue

        assert os.path.exists(emd_of_img_ori_path),emd_of_img_ori_path
        assert os.path.exists(img_of_img_edit_path),img_of_img_edit_path
        assert os.path.exists(mask_of_img_mask_path),mask_of_img_mask_path

        concat_res,res_image = generate()
        save_name = 'tmp.jpg'
        concat_res.save(save_name)

        # pdb.set_trace()

        res_image.save(save_path)



def test_holes():
    
    examples_dir = '/data/shengjie/style_zhenzhi/'
    save_dir = '/data/shengjie/synthesis_zhenzhi/'

    imagefiles = os.listdir(examples_dir)

    test_img = os.path.join(examples_dir,imagefiles[0])
    # PIL
    image = load_image(test_img) # PIL
    image,_,_,_,_ = pad_image(image)
    image = image.resize(target_shape)

    mask_ori,mask_pil_ori = get_mask_by_sam(test_img)
    for hole_size,hole_nums in product([50,100],[10,15,20,25,30]):
        save_name = f'./tmp_hsize{hole_size}_hnums{hole_nums}.png'
        mask_hole , mask_hole_pil = add_random_holes(mask_pil_ori,
                                                    min_holes=hole_nums,
                                                    max_holes=hole_nums,
                                                    hole_size_range=(hole_size,hole_size))
        # pdb.set_trace()
        mask_pil = mask_hole_pil
        mask,_,_,_,_ = pad_image(mask_pil,fill=(255,))
        mask = mask.resize(target_shape)

        # pdb.set_trace()
            
        res_image = pipe(
            # prompt="a white paper cup",
            image=image,
            mask_image=mask,
            height=target_shape[1],
            width=target_shape[0],
            guidance_scale=30,
            num_inference_steps=8,
            max_sequence_length=512,
            # generator=torch.Generator("cpu").manual_seed(0),
            **main_condition_prompt,
        ).images[0]

        concat_res = horizontal_concat_images([image,mask,res_image])
        concat_res.save(save_name)

    # pdb.set_trace()


# res_image.save(f"flux-fill-dev.png")